import os

from skimage import io
from torch.utils.data import Dataset
from torchvision import transforms


class CustomDataset(Dataset):
    def __init__(self, input_dir, mask_dir, transform=None):
        self.input_dir = input_dir
        self.mask_dir = mask_dir
        self.input_name = os.listdir(input_dir)
        self.transform = transform

    def __len__(self):
        return len(self.input_name)

    def __getitem__(self, idx):
        img_path = os.path.join(self.input_dir, self.input_name[idx])
        mask_path = os.path.join(self.mask_dir, self.input_name[idx])

        image = io.imread(img_path)
        mask = io.imread(mask_path, as_gray=True)
        mask = mask.squeeze(0)  # 先去掉大小为1的维度
        mask = mask.reshape(image.shape[0], image.shape[1], 1)  # 在最后一个维度添加一个大小为1的维度

        # 调整图像大小
        image = transforms.Resize((224, 224))(transforms.ToTensor()(image))
        mask = transforms.Resize((224, 224))(transforms.ToTensor()(mask))

        # 应用数据转换
        if self.transform:
            image = self.transform(image)

        return image, mask
